#include "main/utils.h"
#include "lib/binary_tree.h"
using namespace std;

BinaryTreeNode* flattenBST(BinaryTreeNode* root) {
  if (!root)
    return nullptr;

  stack<BinaryTreeNode*> help_stack;
  BinaryTreeNode* cur_node = root;
  BinaryTreeNode* pre_node = nullptr;
  BinaryTreeNode* new_root = nullptr;
  while (cur_node || !help_stack.empty()) {
    while (cur_node) {
      help_stack.push(cur_node);
      cur_node = cur_node->left;
    }
    cur_node = help_stack.top();
    help_stack.pop();
    cur_node->left = nullptr;
    if (pre_node) {
      pre_node->right = cur_node;
    } else {
      new_root = cur_node;
    }
    pre_node = cur_node;
    cur_node = cur_node->right;
  }
  return new_root;
}

int main() {
  vector<int> vec = {4, 2, 5, 1, 3, 5};
  BinaryTreeNode* root = CreateBinaryTree(vec);
  PrintBinaryTree(root);
  auto new_root = flattenBST(root);
  PrintBinaryTree(new_root);

  return 0;
}
